
import json
import pydot
import os
import logging
import sys,subprocess
import argparse
import re,utils


def create_dot_from_json(data,dot_dir):
    graph = pydot.Dot(graph_type='digraph')
    graph_with_explain = pydot.Dot(graph_type='digraph')
    state_machine_names = data.keys()
    for state_machine_name in state_machine_names:
        dot_path = os.path.join(dot_dir,f"{state_machine_name}.dot")
        # print(data[state_machine_name])
        # print(type(data[state_machine_name]))
        for key, value in data[state_machine_name].items():
            transitions = value['transitions']


            for transition in transitions:
                start_state = transition["start_state"]
                end_state = transition["end_state"]
                condition = transition["condition"]        
                explanation = transition['condition explanation']

                edge = pydot.Edge(start_state, end_state, label=condition)
                graph.add_edge(edge)

        graph.write(dot_path)
        pdf_path = dot_path[:-3] + 'pdf'
        command = ['dot','-Tpdf',dot_path,'-o',pdf_path]

        subprocess.run(command)


def main(input,output):
    os.makedirs(output,exist_ok=True)
    with open(input, 'r', encoding='utf-8') as f:
        json_data = json.load(f)
    
    create_dot_from_json(json_data, os.path.abspath(output))

# /bin/python3 /home/ju/Desktop/llmstate/json2fsm.py -i ./save/openssl_3.4.1/llm_fsm_response.json -o ./save/openssl_3.4.1/dots

# python ./json2fsm.py -i ./save/openssl_3.4.1/get_fsm_from_llm.json -o ./save/openssl_3.4.1/dots
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='json to fsm')
    parser.add_argument('-i', '--input', help='Input config JSON file path', required=True)
    parser.add_argument('-o', '--output', help='Output dir', required=True)

    args = parser.parse_args()


    main(args.input,args.output)

    